package com.hyq.learn.ai.filter;

import jakarta.servlet.ServletOutputStream;
import jakarta.servlet.WriteListener;
import jakarta.servlet.http.HttpServletResponse;
import jakarta.servlet.http.HttpServletResponseWrapper;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.nio.charset.StandardCharsets;

/**
 * @author heyuqin
 * @date 2025/11/4 15:33
 */
public class LoggingSseResponseWrapper extends HttpServletResponseWrapper {

    public LoggingSseResponseWrapper(HttpServletResponse response) {
        super(response);
    }

    @Override
    public ServletOutputStream getOutputStream() throws IOException {
        return new LoggingServletOutputStream(super.getOutputStream());
    }

    public static class LoggingServletOutputStream extends ServletOutputStream {

        private static final Logger LOGGER = LoggerFactory.getLogger(LoggingServletOutputStream.class);

        private final ServletOutputStream original;
        private final ByteArrayOutputStream buffer = new ByteArrayOutputStream();

        public LoggingServletOutputStream(ServletOutputStream original) {
            this.original = original;
        }

        @Override
        public void write(int b) throws IOException {
            original.write(b);
            buffer.write(b);
        }

        @Override
        public void write(byte[] b, int off, int len) throws IOException {
            original.write(b, off, len);
            buffer.write(b, off, len);
        }

        @Override
        public void flush() throws IOException {
            original.flush();
            // flush 时打印日志
            logBuffer();
            buffer.reset();
        }

        @Override
        public void close() throws IOException {
            original.close();
            // close 时打印剩余缓存
            logBuffer();
            buffer.reset();
        }

        private void logBuffer() {
            if (buffer.size() > 0) {
                String s = buffer.toString(StandardCharsets.UTF_8);
                LOGGER.info("SSE chunk completed: {}", s);
            }
        }

        @Override
        public boolean isReady() {
            return original.isReady();
        }

        @Override
        public void setWriteListener(WriteListener listener) {
            original.setWriteListener(listener);
        }
    }
}
